import pickle

import numpy as np
import scipy.special


class neuralNetwork:
    # 初始化神经网络
    def __init__(self, inputnodes, hiddennodes, outputnodes, learningrate):
        # 设置输入层、隐藏层、输出层节点的数量
        self.inodes = inputnodes
        self.hnodes = hiddennodes
        self.onodes = outputnodes

        # 连接权重，随机生成输入层到隐藏层和隐藏层到输出层的权重
        self.wih = np.random.rand(self.hnodes, self.inodes) - 0.5
        self.who = np.random.rand(self.onodes, self.hnodes) - 0.5

        # 学习率
        self.lr = learningrate
        pass

    # 将激活函数设置为 sigmoid 函数
    def activation_function(self, x):
        return scipy.special.expit(x)

    # 训练神经网络
    def train(self, inputs_list, targets_list):
        # 将输入数据与目标数据转为二维数组
        inputs = np.array(inputs_list, ndmin=2).T
        targets = np.array(targets_list, ndmin=2).T

        # 通过矩阵点积和激活函数得到隐藏层的输出
        hidden_inputs = np.dot(self.wih, inputs)
        hidden_outputs = self.activation_function(hidden_inputs)

        # 通过矩阵点积和激活函数得到最终输出
        final_inputs = np.dot(self.who, hidden_outputs)
        final_outputs = self.activation_function(final_inputs)

        # 获取目标值与实际值的差值
        output_errors = targets - final_outputs
        # 反向传播差值
        hidden_errors = np.dot(self.who.T, output_errors)

        # 通过梯度下降法更新隐藏层到输出层的权重
        self.who += self.lr * np.dot(
            (output_errors * final_outputs * (1.0 - final_outputs)),
            np.transpose(hidden_outputs)
        )
        # 通过梯度下降法更新输入层到隐藏层的权重
        self.wih += self.lr * np.dot(
            (hidden_errors * hidden_outputs * (1.0 - hidden_outputs)),
            np.transpose(inputs)
        )

        pass

    # query the neural network
    def query(self, inputs_list):
        # 将输入的数组转化为一个二维数组
        inputs = np.array(inputs_list, ndmin=2).T

        # 计算输入数据与权重的点积
        hidden_inputs = np.dot(self.wih, inputs)
        # 经过激活函数的到隐藏层数据
        hidden_outputs = self.activation_function(hidden_inputs)

        # 计算隐藏层数据与权重的点积
        final_inputs = np.dot(self.who, hidden_outputs)
        # 最终到达输出层的数据
        final_outputs = self.activation_function(final_inputs)

        return final_outputs

    def saveModel(self):
        modelDict = {"wih": self.wih, "who": self.who, "lr": self.lr}
        for val in modelDict.keys():
            print(val+":")
            print(modelDict[val], "\n")
        a_file = open("data.pkl", "wb")
        pickle.dump(self, a_file)
        a_file.close()

    # 神经网络类定义
    def loadModel(self):
        a_file = open("data.pkl", "rb")
        model = pickle.load(a_file)
        a_file.close()
        return model


# 设置每一层的节点数量
input_nodes = 784
hidden_nodes = 100
output_nodes = 10

# 学习率
learning_rate = 0.3

# 创建神经网络模型
n = neuralNetwork(input_nodes, hidden_nodes, output_nodes, learning_rate)

# 加载训练数据
training_data_file = open("datasets/mnist_train.csv", 'r')
training_data_list = training_data_file.readlines()
training_data_file.close()

# 训练神经网络
# epochs 表示训练次数
epochs = 10
for e in range(epochs):
    # 遍历所有数据进行训练
    for record in training_data_list:
        # 数据通过 ',' 分割，变成一个数组
        all_values = record.split(',')
        # 分离出图片的像素点到一个单独数组
        inputs = (np.asfarray(all_values[1:]) / 255.0 * 0.99) + 0.01
        # 创建目标输出值（数字 0~9 出现的概率，默认全部为 0.01）
        targets = np.zeros(output_nodes) + 0.01
        # all_values[0] 表示手写数字的真实值，将该数字的概率设为 0.99
        targets[int(all_values[0])] = 0.99
        n.train(inputs, targets)
        pass
    pass

# 训练完毕
print('done')

# 加载测试数据
test_data_file = open("datasets/mnist_test.csv", 'r')
test_data_list = test_data_file.readlines()
test_data_file.close()

# 测试神经网络
# 记录所有的训练值，正确存 1 ，错误存 0 。
scorecard = []

# 遍历所有数据进行测试
for record in test_data_list:
    # 数据通过 ',' 分割，变成一个数组
    all_values = record.split(',')
    # 第一个数字为正确答案
    correct_label = int(all_values[0])
    # 取出测试的输入数据
    inputs = (np.asfarray(all_values[1:]) / 255.0 * 0.99) + 0.01
    # 查询神经网络
    outputs = n.query(inputs)
    # 取出概率最大的数字，表示输出
    label = np.argmax(outputs)
    # 打印出真实值与查询值
    print('act: ', label, ' pre: ', correct_label)
    if label == correct_label:
        # 神经网络查询结果与真实值匹配，记录数组存入 1
        scorecard.append(1)
    else:
        # 神经网络查询结果与真实值不匹配，记录数组存入 0
        scorecard.append(0)
        pass

    pass

# 计算训练的成功率
scorecard_array = np.asarray(scorecard)
print("performance = ", scorecard_array.sum() / scorecard_array.size)
n.saveModel()
model = n.loadModel()
